
import numpy as np
import onnx
import onnx_graphsurgeon as gs
from copy import deepcopy

src_onnx = "onnx/clip_opt.onnx"
dst_onnx = "clip_opt.onnx"

graph = gs.import_onnx(onnx.shape_inference.infer_shapes(onnx.load(src_onnx)))

for node in graph.nodes:
    if node.op == "Gather":
        node.op = "GatherV2"
        continue
    if node.op == "Softmax" and node.i().op == "Reshape":
        add_node = node.i().i()
        w = add_node.inputs[1].values
        w[w == -np.inf] = -10000

graph.cleanup()
onnx.save(gs.export_onnx(graph), dst_onnx)
